#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdlib.h>

#define _USE_MATH_DEFINES
#include <math.h>

#include "raylib.h"
#include "pub_lib.h"
#include "FreeImage.h"
#include "image.h"
#include "dpi_algo.h"

void imgage_filter_common(FIBITMAP* ori, FIBITMAP* chg, int dim,\
	spatial_filter_func filter, void* extra)
{
	unsigned int x, y, p, q, k, index;
	BYTE value, new_value;

	unsigned int w, h;
	w = FreeImage_GetWidth(ori);
	h = FreeImage_GetHeight(ori);

	unsigned int dis = dim / 2;
	int sum = 0;
	new_value = 0;
	k = 0; //0 based index of mask[dim][dim]
	for(y = dis; y < (h - dis); y++){
		for(x = dis; x < (w - dis); x++) {
			k = 0;
			index = y * w + x;
			for(q = y - dis; q <= y + dis; q++) {
				for(p = x - dis; p <= x + dis; p++) {
					k++;
					FreeImage_GetPixelIndex(ori, p, q, &value);
					filter(index, &sum, k, dim, value, &new_value, extra);
				}//end for q
			}//end for p
			FreeImage_SetPixelIndex(chg, x, y, &new_value);
		}//end for x
	} //end for y
}

///////////////////////////////////////////////////////////////////////////////

/*
 * 均值滤波
 * dim = 3, 3 * 3 averaging mask
 * dim = 5, 5 * 5 averaging mask
 * */
void spatial_filter_avg(unsigned int index, int* sum, unsigned char k,\
	int dim, BYTE val, BYTE* new_val, void* extra)
{
	int end = dim * dim;
	if (k == 1)
		*sum = val;
	else
		*sum += val;

	if (k == end)
		*new_val = *sum / end;
}

void image_filter_avg(FIBITMAP* ori, FIBITMAP* chg, int dim)
{
	imgage_filter_common(ori, chg, dim, spatial_filter_avg, NULL);
}

///////////////////////////////////////////////////////////////////////////////

/* 中值滤波 */

typedef struct{
	sp_filter_stat type;
	BYTE* data;
}sp_filter_stat_str;

static void spatial_filter_stat_common(unsigned int index, int* sum,
	unsigned char k, int dim, BYTE val, BYTE* new_val, void* extra)
{
	int end = dim * dim;
	int i, ret;
	BYTE* data;
	sp_filter_stat_str* pstat;

	pstat = (sp_filter_stat_str*)extra;
	data = pstat->data;
	data[k - 1] = val;

	if (k == end) {
		switch(pstat->type){
			case SPSTAT_MEDIAN:
				i = (end + 1) / 2;
				break;
			case SPSTAT_MAX:
				i = end;
				break;
			case SPSTAT_MIN:
			default:
				i = 1;
				break;
		}

		ret = kthSmallest(data, 0, end - 1, i, ascending);
		*new_val = data[ret];
	}
}

void image_filter_stat_common(FIBITMAP* ori, FIBITMAP* chg,
	int dim, sp_filter_stat type)
{
	int size;
	size = dim * dim;

	sp_filter_stat_str stat;
	stat.type = type;
	stat.data = calloc(size, 1);
	imgage_filter_common(ori, chg, dim, spatial_filter_stat_common, (void*)&stat);

	free(stat.data);
}


///////////////////////////////////////////////////////////////////////////////
//3.6 锐化空间滤波器
//laplacian

//3.37 a,b, C = -1, ori - mask
char laplacian_coff1[9] = {0, 1, 0, 1, -4, 1, 0, 1, 0};
char laplacian_coff2[9] = {1, 1, 1, 1, -8, 1, 1, 1, 1};

//3.37 c,d, C = 1, ori + mask
char laplacian_coff3[9] = {0, -1, 0, -1, 4, -1, 0, -1, 0};
char laplacian_coff4[9] = {-1, -1, -1, -1, 8, -1, -1, -1, -1};

typedef struct{
	char* coff;
	float* data;
}laplacian_data;

void spatial_filter_laplacian(unsigned int index, int* sum, unsigned char k,\
	int dim, BYTE val, BYTE* new_val, void* extra)
{
	int end = dim * dim;
	laplacian_data* lap = (laplacian_data*)extra;
	char* coff = lap->coff;
	float* data = lap->data;

	if (k == 1)
		*sum = val * coff[k - 1];
	else
		*sum += (val * coff[k - 1]);

	if (k == end) {
		if(data)
			data[index] = *sum;

		if(*sum < 0)
			*sum = 0;
		if(*sum > 255)
			*sum = 255;
		*new_val = *sum;
	}
}

/*
 * type = 1, used laplacian_coff1
 * type = 2, used laplacian_coff1
 * c = -1
 * otherwise c = 1
 * data[h][w]: 存放中间计算结果，因为计算的结果可能是负值
 * output:
 *     chg = laplacian result
 * */
void image_filter_lapaction(FIBITMAP* ori, FIBITMAP* chg, int type, float* data)
{
	char* extra;
	switch (type) {
	case 1:
		extra = laplacian_coff1;
		break;
	case 2:
		extra = laplacian_coff2;
		break;
	case 3:
		extra = laplacian_coff3;
		break;
	case 4:
	default:
		extra = laplacian_coff4;
		break;
	}

	laplacian_data lap;
	lap.coff = extra;
	lap.data = data;
	if(data)
		init_image_data(ori, data);
	imgage_filter_common(ori, chg, 3, spatial_filter_laplacian, (void*)&lap);
}

//chg = ori + c * laplacian result
void image_filter_lapaction_final(FIBITMAP* ori, FIBITMAP* chg, int type, float* data)
{
	int c;
	image_filter_lapaction(ori, chg, type, data);
	if(type == 1 || type == 2)
		c = -1;
	else
		c = 1;
	image_add(ori, chg, chg, c);
}

/*
 * 空间滤波器 mask
 * char mask[dim][dim]
 * */
void image_filter_mask(FIBITMAP* ori, FIBITMAP* chg, int dim, char* mask, float* data)
{
	laplacian_data lap;
	lap.coff = mask;
	lap.data = data;
	if(data)
		init_image_data(ori, data);
	imgage_filter_common(ori, chg, 3, spatial_filter_laplacian, (void*)&lap);
}

///////////////////////////////////////////////////////////////////////////////
/*
 * Gaussian filter
 * h(x,y) = exp(-(x^2 + y^2) / 2sigma^2)
 * */
void product_gaussian_coef(int dim, float sigma, float* coef)
{
	double  div, exponent;
	int x, y, dis, sum, k;
	float total = 0.0f;

	dis = dim / 2;
	k = 0;
	for(y = -dis; y <= dis; y++) {
		for(x = -dis; x <= dis; x++) {
			sum = x * x + y * y;
			div = sigma * sigma * (-2.0);
			exponent = sum / div;
			coef[k] = exp(exponent);
			total += coef[k];
			k++;
		}
	}

	int size = dim * dim;
	for (x = 0; x < size; x++)
		coef[x] = coef[x] / total;
}

void spatial_filter_gaussian(unsigned int index, int* sum, unsigned char k,\
	int dim, BYTE val, BYTE* new_val, void* extra)
{
	int end = dim * dim;
	float *data;
	float temp, pre_sum = 0.0f;
	data = (float*)extra;

	temp = val * data[k - 1];

	if (k == 1) {
		memcpy(sum, &temp, 4);
	} else {
		memcpy(&pre_sum, sum, 4);
		pre_sum += temp;
		memcpy(sum, &pre_sum, 4);
	}

	if (k == end) {
		*new_val = pre_sum;
	}
}

void print_coef(int dim, float* coef, char* name)
{
	int i, j, k = 0;
	float total = 0.0f;

	if(name)
		printf("%s %d * %d coefficient: \n", name, dim, dim);
	else
		printf("%d * %d coefficient: \n", dim, dim);

	for (i = 0; i < dim; i++) {
		for (j = 0; j < dim; j++) {
			printf("%.4f ", coef[k]);
			total += coef[k];
			k++;
		}
		printf("\n");
	}
	printf("coeff sum = %.4f\n", total);
}

void image_filter_gaussian(FIBITMAP* ori, FIBITMAP* chg, int dim, float sigma)
{
	int size;
	float *coef;

	size = dim * dim;
	coef = malloc(size * sizeof(float));
	//1. product gaussian filter coefficient
	product_gaussian_coef(dim, sigma, coef);
	print_coef(dim, coef, "Gaussian");
	//2. blur the original image, filter
	imgage_filter_common(ori, chg, dim, spatial_filter_gaussian, (void*)coef);
	free(coef);
}

///////////////////////////////////////////////////////////////////////////////
//filter: sobel, grad
//3.6.4 梯度图，Sobel 算子

typedef struct{
	int ysum;
	int* gx;
	int* gy;
}sobel_data;

//figure 3.41, d, e
char sobel_coffx[9] = {-1, -2, -1, 0, 0, 0, 1, 2, 1};
char sobel_coffy[9] = {-1, 0, 1, -2, 0, 2, -1, 0, 1};

/* 3 * 3
 * sum:  sum x
 * extra: sum y
 * */
void spatial_filter_grad_sobel(unsigned int index, int* sum, unsigned char k,\
	int dim, BYTE val, BYTE* new_val, void* extra)
{
	int end = dim * dim;
	int x, y, temp;
	sobel_data* psobel = (sobel_data*)extra;
	int *ysum = &psobel->ysum;

	x = val * sobel_coffx[k - 1];
	y = val * sobel_coffy[k - 1];

	if (k == 1) {
		*sum = x;
		*ysum = y;
	} else {
		*sum += x;
		*ysum += y;
	}

	if (k == end) {
		x = *sum;
		y = *ysum;
		if(psobel->gx)
			psobel->gx[index] = x;
		if(psobel->gy)
			psobel->gy[index] = y;

		temp = abs(x) + abs(y);
		if(temp > 255)
			temp = 255;
		*new_val = temp;
	}
}

void image_filter_grad_sobel_ext(FIBITMAP* ori, FIBITMAP* chg, int* gx, int* gy)
{
	sobel_data sobel;
	sobel.gx = gx;
	sobel.gy = gy;
	if(gx)
		init_image_datai(ori, gx);
	if(gy)
		init_image_datai(ori, gy);
	imgage_filter_common(ori, chg, 3, spatial_filter_grad_sobel, &sobel);
}

///////////////////////////////////////////////////////////////////////////////

static float calc_ftheta(int gxx, int gyy, int gxy, float theta)
{
	float a, b, c;

	a = gxx + gxy;
	b = (gxx - gyy) * cosf(theta);
	c = 2 * gxy * sinf(theta);
	return  (a + b + c);
}

/*
 * 6.7-3 - 6.7-9
 * calculate the gradient(magnitude and direction) in RGB space
 * result saved in float data[h][w]
 * */
void color_image_gradient(FIBITMAP* color, float* data)
{
	int w, h;
	w = FreeImage_GetWidth(color);
	h = FreeImage_GetHeight(color);
	int size = w * h * sizeof(int);

	int *rgx = malloc(size);
	int *rgy = malloc(size);

	int *ggx = malloc(size);
	int *ggy = malloc(size);

	int *bgx = malloc(size);
	int *bgy = malloc(size);

	FIBITMAP* ori = FreeImage_ConvertTo8Bits(color);
	FIBITMAP* chg = copy_img(ori);

	init_image_data(ori, data); //init data[]
	//rx, ry
	fill_image_data_color(color, ori, 0);
	image_filter_grad_sobel_ext(ori, chg, rgx, rgy);

	//gx, gy
	fill_image_data_color(color, ori, 1);
	image_filter_grad_sobel_ext(ori, chg, ggx, ggy);

	//bx, by
	fill_image_data_color(color, ori, 2);
	image_filter_grad_sobel_ext(ori, chg, bgx, bgy);

	int rx, ry, gx, gy, bx, by;
	int gxx, gyy, gxy;

	//fill data[]
	int x, y, index;
	int dis = 1; //3 * 3 mask, so dis = 1
	float temp, theta, a, b, c;
	for(y = dis; y < (h - dis); y++){
		for(x = dis; x < (w - dis); x++) {
			index = y * w + x;
			rx = rgx[index];
			ry = rgy[index];

			gx = ggx[index];
			gy = ggy[index];

			bx = bgx[index];
			by = bgy[index];

			gxx = rx * rx + gx * gx + bx * bx;
			gyy = ry * ry + gy * gy + by * by;
			gxy = rx * ry + gx * gy + bx * by;

			temp = (2 * gxy) / (gxx - gyy + MIN_NUMBER);
			theta = atanf(temp);

			a = calc_ftheta(theta, gxx, gyy, gxy);
			b = calc_ftheta(theta + M_PI, gxx, gyy, gxy);
			c = max(a, b) * 0.5;
			if(c < 0)
				c = 0 - c;
			data[index] = sqrtf(c);
		}//end for x
	} //end for y

	close_img(ori);
	close_img(chg);
	free(rgx);
	free(rgy);
	free(ggx);
	free(ggy);
	free(bgx);
	free(bgy);
}

//calc RGB gradient, and merge to data
void color_image_gradient_merge(FIBITMAP* color, float* data)
{
	FIBITMAP* ori = FreeImage_ConvertTo8Bits(color);
	FIBITMAP* chg = copy_img(ori);

	int w, h;
	w = FreeImage_GetWidth(color);
	h = FreeImage_GetHeight(color);
	int size = w * h;
	int fsize = size * sizeof(float);
	float *image_data = malloc(fsize);

	int i, j;
	for(i = 0; i < size; i++)
		data[i] = 0;

	for(i = 0; i < 3; i++) {
		fill_image_data_color(color, ori, i); //fill certain color
		image_filter_grad_sobel(ori, chg); //filter
		init_image_data(chg, image_data); //write to image_data
		//add RGB components
		for(j = 0; j < size; j++)
			data[j] += image_data[j];
	}

	free(image_data);
	close_img(ori);
	close_img(chg);
}

///////////////////////////////////////////////////////////////////////////////


///////////////////////////////////////////////////////////////////////////////


///////////////////////////////////////////////////////////////////////////////

/* 图像标定
 * 让灰度值分布均匀
value = 255 * (v - min) / max
*/
void image_scale(FIBITMAP* dib, BYTE min, BYTE max)
{
	//val = 255 * (val - min) / max
	int x, y;
	BYTE value;

	int w, h;
	w = FreeImage_GetWidth(dib);
	h = FreeImage_GetHeight(dib);
	for (y = 0; y < h; y++) {
		for (x = 0; x < w; x++) {
			FreeImage_GetPixelIndex(dib, x, y, &value);
			value = 255 * (value - min) / max; //标定
			FreeImage_SetPixelIndex(dib, x, y, &value);
		}
	}
}

/*
 * 标定一副图片
 * 标定到 [0, 255]
 * */
void image_filter_scale_show(FIBITMAP* dib, float* data)
{
	BYTE* out;
	int w, h;
	w = FreeImage_GetWidth(dib);
	h = FreeImage_GetHeight(dib);

	out = malloc(w * h);
	scale_float_data(data, out, w * h);
	fill_image_data(dib, out);

	free(out);
}

void image_filter_scale(FIBITMAP* dib)
{
	BYTE min, max;
	image_get_intensity_range(dib, &min, &max);
	image_scale(dib, min, max);
}

///////////////////////////////////////////////////////////////////////////////




///////////////////////////////////////////////////////////////////////////////


int ascending(BYTE x, BYTE y)
{
	return x <= y;
}

int descending(BYTE x, BYTE y)
{
	return x >= y;
}

// Standard partition process of QuickSort().
// It considers the last element as pivot
// and moves all smaller element to left of
// it and greater elements to right
int partition(BYTE* arr, int l, int r, func_compare pf)
{
	int x = arr[r], i = l;
	for (int j = l; j <= r - 1; j++) {
		//if (arr[j] <= x) {
		if(pf(arr[j], x)) {
			swapB(arr[i], arr[j]);
			i++;
		}
	}
	swapB(arr[i], arr[r]);
	return i;
}

// This function returns k'th smallest
// element in arr[l..r] using QuickSort
// based method. ASSUMPTION: ALL ELEMENTS
// IN ARR[] ARE DISTINCT
int kthSmallest(BYTE* arr, int l, int r, int k, func_compare pf)
{
	// If k is smaller than number of
	// elements in array
	if (k > 0 && k <= r - l + 1) {

		// Partition the array around last
		// element and get position of pivot
		// element in sorted array
		int index = partition(arr, l, r, pf);

		// If position is same as k
		if (index - l == k - 1)
			return index;

		// If position is more, recur
		// for left subarray
		if (index - l > k - 1)
			return kthSmallest(arr, l, index - 1, k, pf);

		// Else recur for right subarray
		return kthSmallest(arr, index + 1, r, k - index + l - 1, pf);
	}

	// If k is more than number of elements in array
	return 0; //error
}

///////////////////////////////////////////////////////////////////////////////


///////////////////////////////////////////////////////////////////////////////



